# -*- coding: utf-8 -*-
"""
Created on Thu Aug 12 16:41:57 2021

@author: weixifei
"""
from PIL import Image
import os
import numpy as np
import cv2
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import cohen_kappa_score

import torchvision.transforms as transforms

from torch.utils.data import DataLoader

class GAMMA_sub1_dataset():
    """
    getitem() output:
        fundus_img: RGB uint8 image with shape (3, image_size, image_size)
        oct_img:    Uint8 image with shape (256, oct_img_size[0], oct_img_size[1])
    """
    # print("GAMMA_sub1_dataset:start")
    def __init__(self,
                 img_transforms,
                 oct_transforms,
                 dataset_root,
                 label_file='',
                 filelists=None,
                 num_classes=3,
                 mode = 'train'):

        self.dataset_root = dataset_root

        # print("__init__")
        self.img_transforms = img_transforms
        self.oct_transforms = oct_transforms
        self.mode = mode.lower()
        self.num_classes = num_classes
        label = {row['data']: row[1:].values for _, row in pd.read_excel(label_file).iterrows()}
        # print(os.listdir(dataset_root))

        self.file_list = [[f, label[int(f)]] for f in os.listdir(dataset_root)]

        train_filelists, val_filelists = train_test_split(self.file_list, test_size=0.2, random_state=42)
        if self.mode == 'train':
            self.file_list = train_filelists
            
        elif self.mode == 'val':
            self.file_list =val_filelists
            
        elif self.mode == "test":
            self.file_list = [[f, None] for f in os.listdir(dataset_root)]
        
        if filelists is not None:   # filelists=None
            self.file_list = [item for item in self.file_list if item[0] in filelists]

    def __getitem__(self, idx):
        # print("__getitem__")
        real_index, label = self.file_list[idx]

        # 得到眼底彩照的路径
        fundus_img_path = os.path.join(self.dataset_root, real_index, real_index + ".jpg")

        # 得到OCT_arr的路径
        oct_series_list = sorted(os.listdir(os.path.join(self.dataset_root, real_index, real_index)), 
                                    key=lambda x: int(x.split("_")[0]))
        # print(oct_series_list)

        fundus_img = cv2.imread(fundus_img_path)[:, :, ::-1]    # BGR -> RGB
        # fundus_img = cv2.imread(fundus_img_path)

        oct_series_0 = cv2.resize(cv2.imread(os.path.join(self.dataset_root, real_index, real_index, oct_series_list[0]), 
                                    cv2.IMREAD_GRAYSCALE),(256,256))
        # 读取 oct_series_list[0] 为了得到size

        oct_img = np.zeros((len(oct_series_list), oct_series_0.shape[0], oct_series_0.shape[1], 1), dtype="uint8")

        for k, p in enumerate(oct_series_list):
            oct_img[k] =cv2.resize(cv2.imread(
                os.path.join(self.dataset_root, real_index, real_index, p), cv2.IMREAD_GRAYSCALE),(256,256))[..., np.newaxis]

        oct_img = oct_img.squeeze(-1)
        if self.img_transforms is not None:
            fundus_img = self.img_transforms(fundus_img)

        if self.oct_transforms is not None:
            oct_img = self.oct_transforms(oct_img)

        oct_img = oct_img.transpose(1, 0) # H, W, C -> C, H, W
        
        if self.mode == 'test':
            return fundus_img, oct_img, real_index
        if self.mode == "train" or 'val':
            # print("train or val",fundus_img.shape,oct_img.shape,label.shape)
            label = label.argmax()
            return fundus_img, oct_img, label

    def __len__(self):
        return len(self.file_list)

    # print("GAMMA_sub1_dataset:end")


if __name__ == "__main":

    image_size = 512
    img_train_transforms = transforms.Compose([
        transforms.ToPILImage(),
        transforms.RandomResizedCrop(image_size, scale=(0.8, 1)),
        transforms.Resize((image_size, image_size)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(15),
        transforms.ColorJitter(contrast=0.4,brightness=0.4),

        # transforms.ColorJitter(contrast=0.4,brightness=0.4),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]),
    ])

    oct_train_transforms = transforms.Compose([
    #        transforms.ToPILImage(),
    #        transforms.Resize((image_size, image_size)),
    #        transforms.RandomHorizontalFlip(),
    #        transforms.RandomRotation(15),
    #        transforms.ColorJitter(contrast=0.2,brightness=0.2),
        transforms.ToTensor(),
    #        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]),
    ])
    trainset_root='../GAMMA_training data/training_data/multi-modality_images'

    GA = GAMMA_sub1_dataset(dataset_root=trainset_root, # 训练数据和val数据在文件夹进行分配
                            img_transforms=img_train_transforms,
                            oct_transforms=oct_train_transforms,
                            label_file='../GAMMA_training data/training_data/glaucoma_grading_training_GT.xlsx',
                            mode='train')
    for i in GA:
        print(i)
